package org.mat.framework.core.context;

import java.util.EmptyStackException;
import java.util.HashMap;
import java.util.Map;

/**
 * @ClassName: RequestContextUtils
 * @Date: 2021/8/17
 * @author: sunxinhe
 * @Version: 1.0
 * @Description: TODO 结构有点乱了，抽空整理下，特别是初始化逻辑，不严谨
 */
public class RequestContextUtils {

    private static final ThreadLocal<RequestContext> threadLocal = new ThreadLocal<>();

    public static void set(RequestContext requestContext) {
        if (requestContext == null) {
            return;
        }
        threadLocal.set(requestContext);
    }

    public static RequestContext get() throws EmptyStackException {
        return threadLocal.get();
    }

    public static String getTraceId() throws EmptyStackException {
        String traceId = null;
        RequestContext requestContext = threadLocal.get();
        if (requestContext != null) {
            traceId = requestContext.getTraceId();
        }
        return traceId;
    }

    public static String get(String key) {
        RequestContext requestContext = get();
        if (requestContext == null) {
            init(new DefaultTraceIdAdaptor().getTraceId());
            requestContext = get();
        }
        Map<String, String> data = requestContext.getData();
        if (data == null) {
            requestContext.setData(new HashMap<>());
            data = requestContext.getData();
        }

        return data.get(key);
    }

    public static void put(String key, String value) {
        RequestContext requestContext = get();
        if (requestContext == null) {
            init(new DefaultTraceIdAdaptor().getTraceId());
            requestContext = get();
        }
        Map<String, String> data = requestContext.getData();
        if (data == null) {
            requestContext.setData(new HashMap<>());
            data = requestContext.getData();
        }

        data.put(key, value);
    }

    public static RequestContext getAndClear() throws EmptyStackException {
        RequestContext requestContext = threadLocal.get();
        clear();
        return requestContext;
    }

    public static void clear() {
        threadLocal.remove();
    }

    public static RequestContext init(String traceId, Boolean debugEnabled, Long startTime) {
        RequestContext requestContext;
        if (get() != null) {
            requestContext = get();
        } else {
            requestContext = new RequestContext();
        }

        requestContext.setTraceId(traceId);
        requestContext.setDebugEnabled(debugEnabled);
        requestContext.setStartTime(startTime);

        set(requestContext);

        return requestContext;
    }

    public static RequestContext init(String traceId) {
        RequestContext requestContext;
        if (get() != null) {
            requestContext = get();
        } else {
            requestContext = new RequestContext();
        }

        requestContext.setTraceId(traceId);

        set(requestContext);

        return requestContext;
    }
}
